{
  "cells": [
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "wdeKOEkv1Fe8"
      },
      "source": [
        "##### Copyright 2021 The TensorFlow Authors."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "cellView": "form",
        "id": "SpNWyqewk8fE"
      },
      "outputs": [],
      "source": [
        "#@title Licensed under the Apache License, Version 2.0 (the \"License\");\n",
        "# you may not use this file except in compliance with the License.\n",
        "# You may obtain a copy of the License at\n",
        "#\n",
        "# https://www.apache.org/licenses/LICENSE-2.0\n",
        "#\n",
        "# Unless required by applicable law or agreed to in writing, software\n",
        "# distributed under the License is distributed on an \"AS IS\" BASIS,\n",
        "# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n",
        "# See the License for the specific language governing permissions and\n",
        "# limitations under the License."
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "hPxSVqcvqPHO"
      },
      "source": [
        "# TFX Python function component tutorial\n"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "WdRDkO2wQHUw"
      },
      "source": [
        "Note: We recommend running this tutorial in a Colab notebook, with no setup\n",
        "required!  Just click \"Run in Google Colab\".\n",
        "\n",
        "<div class=\"devsite-table-wrapper\"><table class=\"tfo-notebook-buttons\" align=\"left\">\n",
        "<td><a target=\"_blank\" href=\"https://www.tensorflow.org/tfx/tutorials/tfx/python_function_component\">\n",
        "<img src=\"https://www.tensorflow.org/images/tf_logo_32px.png\" />View on TensorFlow.org</a></td>\n",
        "<td><a target=\"_blank\" href=\"https://colab.research.google.com/github/tensorflow/tfx/blob/master/docs/tutorials/tfx/python_function_component.ipynb\">\n",
        "<img src=\"https://www.tensorflow.org/images/colab_logo_32px.png\">Run in Google Colab</a></td>\n",
        "<td><a target=\"_blank\" href=\"https://github.com/tensorflow/tfx/tree/master/docs/tutorials/tfx/python_function_component.ipynb\">\n",
        "<img width=32px src=\"https://www.tensorflow.org/images/GitHub-Mark-32px.png\">View source on GitHub</a></td>\n",
        "<td><a target=\"_blank\" href=\"https://storage.googleapis.com/tensorflow_docs/tfx/docs/tutorials/tfx/python_function_component.ipynb\">\n",
        "<img width=32px src=\"https://www.tensorflow.org/images/download_logo_32px.png\">Download notebook</a></td>\n",
        "</table></div>"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "GkwEwr0FQKJW"
      },
      "source": [
        "This notebook contains an examples on how to author and run Python function\n",
        "components within the TFX InteractiveContext and in a locally-orchestrated TFX\n",
        "pipeline.\n",
        "\n",
        "For more context and information, see the [Custom Python function components](https://www.tensorflow.org/tfx/guide/custom_function_component)\n",
        "page on the TFX documentation site."
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "d-Av6cm0oBFV"
      },
      "source": [
        "## Setup\n",
        "\n",
        "We will first install TFX and import necessary modules. TFX requires Python 3."
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "fwZ0aXisoBFW"
      },
      "source": [
        "### Check the system Python version\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "PQ-QwavmqPHP"
      },
      "outputs": [],
      "source": [
        "import sys\n",
        "sys.version"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "88XXV3_DQ_Nx"
      },
      "source": [
        "### Upgrade Pip\n",
        "\n",
        "To avoid upgrading Pip in a system when running locally, check to make sure\n",
        "that we're running in Colab.  Local systems can of course be upgraded\n",
        "separately."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "UHvIH-wORCuV"
      },
      "outputs": [],
      "source": [
        "try:\n",
        "  import colab\n",
        "  !pip install --upgrade pip\n",
        "except:\n",
        "  pass"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "FDFSXMLLRGJC"
      },
      "source": [
        "### Install TFX\n",
        "\n",
        "**Note: In Google Colab, because of package updates, the first time you run\n",
        "this cell you must restart the runtime (Runtime > Restart runtime ...).**"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "wGpQOmYIVlSV"
      },
      "outputs": [],
      "source": [
        "!pip install -U tfx"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "hvXNj9a4VlSV"
      },
      "source": [
        "## Did you restart the runtime?\n",
        "\n",
        "If you are using Google Colab, the first time that you run the cell above, you\n",
        "must restart the runtime (Runtime > Restart runtime ...). This is because of\n",
        "the way that Colab loads packages."
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "0v45llP5ReKx"
      },
      "source": [
        "### Import packages\n",
        "We import TFX and check its version."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "bRY0RFJ0VlSV"
      },
      "outputs": [],
      "source": [
        "# Check version\n",
        "from tfx import v1 as tfx\n",
        "tfx.__version__"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "hvb0SspyqPH4"
      },
      "source": [
        "## Custom Python function components\n",
        "\n",
        "In this section, we will create components from Python functions. We will notbe\n",
        "doing any real ML problem — these simple functions are just used to illustrate\n",
        "the Python function component development process.\n",
        "\n",
        "See [Python function based component\n",
        "guide](https://www.tensorflow.org/tfx/guide/custom_function_component)\n",
        "for more documentation."
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "mhYjn9Fj6mdo"
      },
      "source": [
        "### Create Python custom components\n",
        "\n",
        "We begin by writing a function that generate some dummy data. This is written\n",
        "to its own Python module file."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "cHNtKTuiqPH4"
      },
      "outputs": [],
      "source": [
        "%%writefile my_generator.py\n",
        "\n",
        "import os\n",
        "import tensorflow as tf  # Used for writing files.\n",
        "\n",
        "from tfx import v1 as tfx\n",
        "\n",
        "# Non-public APIs, just for showcase.\n",
        "from tfx.types.experimental.simple_artifacts import Dataset\n",
        "\n",
        "@tfx.dsl.components.component\n",
        "def MyGenerator(data: tfx.dsl.components.OutputArtifact[Dataset]):\n",
        "  \"\"\"Create a file with dummy data in the output artifact.\"\"\"\n",
        "  with tf.io.gfile.GFile(os.path.join(data.uri, 'data_file.txt'), 'w') as f:\n",
        "    f.write('Dummy data')\n",
        "\n",
        "  # Set metadata and ensure that it gets passed to downstream components.\n",
        "  data.set_string_custom_property('my_custom_field', 'my_custom_value')"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "MtapXcbSqPH6"
      },
      "source": [
        "Next, we write a second component that uses the dummy data produced.\n",
        "We will just calculate hash of the data and return it."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "27ZEf2xQqPH7"
      },
      "outputs": [],
      "source": [
        "%%writefile my_consumer.py\n",
        "\n",
        "import hashlib\n",
        "import os\n",
        "import tensorflow as tf\n",
        "\n",
        "from tfx import v1 as tfx\n",
        "\n",
        "# Non-public APIs, just for showcase.\n",
        "from tfx.types.experimental.simple_artifacts import Dataset\n",
        "from tfx.types.standard_artifacts import String\n",
        "\n",
        "@tfx.dsl.components.component\n",
        "def MyConsumer(data: tfx.dsl.components.InputArtifact[Dataset],\n",
        "               hash: tfx.dsl.components.OutputArtifact[String],\n",
        "               algorithm: tfx.dsl.components.Parameter[str] = 'sha256'):\n",
        "  \"\"\"Reads the contents of data and calculate.\"\"\"\n",
        "  with tf.io.gfile.GFile(\n",
        "      os.path.join(data.uri, 'data_file.txt'), 'r') as f:\n",
        "    contents = f.read()\n",
        "  h = hashlib.new(algorithm)\n",
        "  h.update(tf.compat.as_bytes(contents))\n",
        "  hash.value = h.hexdigest()\n",
        "\n",
        "  # Read a custom property from the input artifact and set to the output.\n",
        "  custom_value = data.get_string_custom_property('my_custom_field')\n",
        "  hash.set_string_custom_property('input_custom_field', custom_value)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "lIrGHQzFqPII"
      },
      "source": [
        "### Run in-notebook with the InteractiveContext\n",
        "Now, we will demonstrate usage of our new components in the TFX\n",
        "InteractiveContext.\n",
        "\n",
        "For more information on what you can do with the TFX notebook\n",
        "InteractiveContext, see the in-notebook [TFX Keras Component Tutorial](https://www.tensorflow.org/tfx/tutorials/tfx/components_keras)."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "j43snQpRqPII"
      },
      "outputs": [],
      "source": [
        "from my_generator import MyGenerator\n",
        "from my_consumer import MyConsumer"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "_1Rf6FCMSnbM"
      },
      "source": [
        "#### Construct the InteractiveContext"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "dEXGvLLmKviI"
      },
      "outputs": [],
      "source": [
        "# Here, we create an InteractiveContext using default parameters. This will\n",
        "# use a temporary directory with an ephemeral ML Metadata database instance.\n",
        "# To use your own pipeline root or database, the optional properties\n",
        "# `pipeline_root` and `metadata_connection_config` may be passed to\n",
        "# InteractiveContext. Calls to InteractiveContext are no-ops outside of the\n",
        "# notebook.\n",
        "from tfx.orchestration.experimental.interactive.interactive_context import InteractiveContext\n",
        "context = InteractiveContext()"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "XxXVtZzdTWg_"
      },
      "source": [
        "#### Run your component interactively with `context.run()`\n",
        "Next, we run our components interactively within the notebook with\n",
        "`context.run()`. Our consumer component uses the outputs of the generator\n",
        "component."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "kfNmI5qULlSA"
      },
      "outputs": [],
      "source": [
        "generator = MyGenerator()\n",
        "context.run(generator)"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "cRxVZIfFLsL4"
      },
      "outputs": [],
      "source": [
        "consumer = MyConsumer(\n",
        "    data=generator.outputs['data'],\n",
        "    algorithm='md5')\n",
        "context.run(consumer)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "pO_Ggc4QTo0B"
      },
      "source": [
        "After execution, we can inspect the contents of the \"hash\" output artifact of\n",
        "the consumer component on disk."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "h4P3Mx_CT0mP"
      },
      "outputs": [],
      "source": [
        "!tail -v {consumer.outputs['hash'].get()[0].uri}"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "RRosvALKSWe0"
      },
      "source": [
        "That's it, and you've now written and executed your own custom components!"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "uvXKtg17O6mF"
      },
      "source": [
        "### Write a pipeline definition\n",
        "\n",
        "Next, we will author a pipeline using these same components. While using the\n",
        "`InteractiveContext` within a notebook works well for experimentation, defining\n",
        "a pipeline lets you deploy your pipeline on local or remote runners for\n",
        "production usage.\n",
        "\n",
        "Here, we will demonstrate usage of the LocalDagRunner running locally on your\n",
        "machine. For production execution, the Airflow or Kubeflow runners may\n",
        "be more suitable."
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "ft9fbSpnU7C6"
      },
      "source": [
        "#### Construct a pipeline"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "NpkQ805-LyJu"
      },
      "outputs": [],
      "source": [
        "import os\n",
        "import tempfile\n",
        "from tfx import v1 as tfx\n",
        "\n",
        "# Select a persistent TFX root directory to store your output artifacts.\n",
        "# For demonstration purposes only, we use a temporary directory.\n",
        "PIPELINE_ROOT = tempfile.mkdtemp()\n",
        "# Select a pipeline name so that multiple runs of the same logical pipeline\n",
        "# can be grouped.\n",
        "PIPELINE_NAME = \"function-based-pipeline\"\n",
        "# We use a ML Metadata configuration that uses a local SQLite database in\n",
        "# the pipeline root directory. Other backends for ML Metadata are available\n",
        "# for production usage.\n",
        "METADATA_CONNECTION_CONFIG = tfx.orchestration.metadata.sqlite_metadata_connection_config(\n",
        "    os.path.join(PIPELINE_ROOT, 'metadata.sqlite'))\n",
        "\n",
        "def function_based_pipeline():\n",
        "  # Here, we construct our generator and consumer components in the same way.\n",
        "  generator = MyGenerator()\n",
        "  consumer = MyConsumer(\n",
        "      data=generator.outputs['data'],\n",
        "      algorithm='md5')\n",
        "\n",
        "  return tfx.dsl.Pipeline(\n",
        "      pipeline_name=PIPELINE_NAME,\n",
        "      pipeline_root=PIPELINE_ROOT,\n",
        "      components=[generator, consumer],\n",
        "      metadata_connection_config=METADATA_CONNECTION_CONFIG)\n",
        "\n",
        "my_pipeline = function_based_pipeline()"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "Mj-Z3cbFWPbK"
      },
      "source": [
        "#### Run your pipeline with the `LocalDagRunner`"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "PLtGO2PkMQbO"
      },
      "outputs": [],
      "source": [
        "tfx.orchestration.LocalDagRunner().run(my_pipeline)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "Ry4vU3mOWeN1"
      },
      "source": [
        "We can inspect the output artifacts generated by this pipeline execution."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "fyvYTsx8Mp1N"
      },
      "outputs": [],
      "source": [
        "!find {PIPELINE_ROOT}"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "M4CsceadWqHp"
      },
      "source": [
        "You have now written your own custom components and orchestrated their\n",
        "execution on the LocalDagRunner! For next steps, check out additional tutorials\n",
        "and guides on the [TFX website](https://www.tensorflow.org/tfx)."
      ]
    }
  ],
  "metadata": {
    "colab": {
      "collapsed_sections": [
        "wdeKOEkv1Fe8"
      ],
      "name": "python_function_component.ipynb",
      "toc_visible": true
    },
    "kernelspec": {
      "display_name": "Python 3",
      "name": "python3"
    }
  },
  "nbformat": 4,
  "nbformat_minor": 0
}
